from functools import cached_property
import logging
from pathlib import Path
from pydantic import BaseModel
import polars as pl

from .tngs import Result as tNGS, Signal

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("Clinical Pathogen")


class Pathogen(BaseModel):
    taxid: str
    rpm: float
    ctrl_rpm: float
    uniq_reads: int
    genus_ranking: int
    genus_rel_abund: float
    pathogen_type: str
    clnsig: int
    target: str

    @cached_property
    def is_bacteria(self) -> bool:
        return self.pathogen_type in ["细菌", "常见机会致病菌", "定植细菌"]

    @cached_property
    def is_fungi(self) -> bool:
        return self.pathogen_type in ["真菌", "定植真菌"]

    @cached_property
    def is_virus(self) -> bool:
        return self.pathogen_type in ["病毒", "DNA病毒", "RNA病毒", "携带病毒", "意义不明"]

    @cached_property
    def is_special(self) -> bool:
        return self.pathogen_type in ["特殊病原体"]

    def calc_tngs_signal(self) -> Signal:
        signal = Signal.S if self.rpm >= 1000 else (Signal.M if self.rpm >= 200 else Signal.W)
        if signal in [Signal.S, Signal.M] and self.ctrl_rpm > 0 and self.rpm / self.ctrl_rpm < 10:
            signal = Signal.B
        return signal

    def calc_mngs_signal(self, sample_type: str) -> Signal:
        # pylint: disable=R0911,R0912,R0915
        match sample_type:
            case "呼吸道":
                match self.clnsig:
                    case 3:
                        signal = Signal.S
                    case 2:
                        match self.pathogen_type:
                            case "细菌":
                                signal = (Signal.S if self.genus_ranking <= 10 else Signal.M) if self.uniq_reads >= 10 else Signal.W
                            case "特殊" | "病毒":
                                signal = Signal.S if self.uniq_reads >= 10 else Signal.M
                            case "真菌":
                                # TaxID: 42068, 耶氏孢子菌
                                signal = (Signal.S if self.uniq_reads >= 50 else Signal.M) if self.taxid == "42068" else (Signal.S if self.genus_ranking <= 10 else Signal.W)
                            case _:
                                signal = Signal.W
                    case 1:
                        match self.pathogen_type:
                            case "细菌":
                                signal = (Signal.S if self.genus_rel_abund >= 0.5 else Signal.M) if self.uniq_reads >= 10 else Signal.W
                            case "真菌" | "病毒":
                                if self.taxid == "5476":
                                    signal = (Signal.S if self.genus_rel_abund >= 0.5 else Signal.M) if self.uniq_reads >= 10 else Signal.W
                                else:
                                    signal = Signal.S if self.uniq_reads >= 200 else (Signal.M if self.uniq_reads >= 50 else Signal.W)
                        if self.is_bacteria or (self.is_fungi and self.taxid == "5476"):
                            # TaxID: 5476, 白色念珠菌
                            signal = Signal.S if self.uniq_reads >= 10 else Signal.M
                        if (self.is_fungi or self.is_virus) and self.taxid != "5476":
                            # TaxID: 5476, 白色念珠菌
                            signal = Signal.S if self.uniq_reads >= 200 else Signal.M
                    case 0:
                        signal = (Signal.S if self.genus_rel_abund >= 0.5 else Signal.M) if self.uniq_reads >= 10 else Signal.W
                    case _:
                        signal = Signal.W
            case "外周血":
                match self.clnsig:
                    case 3:
                        signal = Signal.S
                    case 2:
                        signal = Signal.S if self.genus_ranking <= 10 else Signal.W
                    case 1:
                        signal = Signal.S
                    case 0:
                        match self.pathogen_type:
                            case "细菌" | "真菌":
                                signal = Signal.S if self.genus_rel_abund >= 0.5 else Signal.M
                            case "病毒":
                                signal = Signal.M if self.genus_rel_abund >= 0.5 else Signal.W
                            case _:
                                signal = Signal.W
                    case _:
                        signal = Signal.W
            case "脑脊液":
                match self.clnsig:
                    case 3:
                        signal = Signal.S
                    case 2:
                        match self.pathogen_type:
                            case "细菌":
                                signal = Signal.S if self.uniq_reads >= 10 and self.genus_ranking <= 5 else Signal.W
                            case "真菌":
                                signal = Signal.S if self.uniq_reads >= 10 else Signal.W
                            case "病毒":
                                signal = Signal.S
                            case _:
                                signal = Signal.W
                    case 1:
                        match self.pathogen_type:
                            case "病毒":
                                signal = Signal.S if self.uniq_reads >= 200 else (Signal.M if self.uniq_reads >= 50 else Signal.W)
                            case _:
                                signal = Signal.W
                    case 0:
                        match self.pathogen_type:
                            case "细菌" | "真菌":
                                signal = Signal.M
                            case _:
                                signal = Signal.W
                    case _:
                        signal = Signal.W
            case "胸水":
                match self.clnsig:
                    case 3:
                        signal = Signal.S
                    case 2:
                        match self.pathogen_type:
                            case "细菌":
                                signal = (Signal.S if self.genus_ranking <= 5 else Signal.M) if self.uniq_reads >= 10 else Signal.W
                            case "特殊" | "病毒":
                                signal = Signal.S if self.uniq_reads >= 10 else Signal.M
                            case "真菌":
                                if self.taxid == "42068":
                                    # TaxID: 42068, 耶氏孢子菌
                                    signal = Signal.S if self.uniq_reads >= 50 else Signal.W
                                else:
                                    signal = Signal.S if self.genus_ranking <= 10 else Signal.W
                            case _:
                                signal = Signal.W
                    case 1:
                        signal = Signal.S if self.uniq_reads >= 200 else Signal.W
                    case 0:
                        match self.pathogen_type:
                            case "细菌" | "真菌":
                                signal = Signal.M
                            case "病毒":
                                signal = (Signal.S if self.genus_rel_abund >= 0.5 else Signal.M) if self.uniq_reads >= 10 else Signal.W
                            case _:
                                signal = Signal.W
                    case _:
                        signal = Signal.W
            case _:
                match self.clnsig:
                    case 3:
                        signal = Signal.S
                    case 2:
                        match self.pathogen_type:
                            case "细菌":
                                signal = (Signal.S if self.genus_ranking <= 5 else Signal.M) if self.uniq_reads >= 10 else Signal.W
                            case "特殊" | "病毒":
                                signal = Signal.S if self.uniq_reads >= 10 else Signal.M
                            case "真菌":
                                signal = Signal.S if self.uniq_reads >= 50 else Signal.W
                            case _:
                                signal = Signal.W
                    case 1:
                        signal = Signal.S if self.uniq_reads >= 200 else (Signal.M if self.uniq_reads >= 50 else Signal.W)
                    case 0:
                        signal = Signal.S if self.genus_rel_abund >= 0.5 else Signal.W
                    case _:
                        signal = Signal.W
        return signal

    def calc_signal(self, sample_type: str) -> Signal:
        return self.calc_tngs_signal() if self.target == "tNGS" else self.calc_mngs_signal(sample_type)


class Result(tNGS):
    clnsig_tsv: list[Path]
    sample_type: str

    @cached_property
    def clnsig_records(self) -> pl.DataFrame:
        records = pl.concat(
            map(
                lambda x: pl.read_csv(
                    x,
                    separator="\t",
                    columns=[0, 3, 4, 5],
                    new_columns=["taxid", "clnsig", "pathogen_type", "sample_type"],
                    schema_overrides={"taxid": str},
                ),
                self.clnsig_tsv,
            )
        )
        records = (
            records.with_columns(pl.col("sample_type").str.split("/").alias("sample_type"))
            .explode(["sample_type"])
            .sort(["clnsig"])
            .filter(pl.col("sample_type").eq(self.sample_type))
            .drop(["sample_type"])
        )
        return records

    @cached_property
    def classified_records(self) -> pl.DataFrame:
        # 注释CLNSIG
        records = self.pathogen_records.join(self.clnsig_records, on="taxid", how="left").with_columns(
            pl.col("clnsig").fill_null(-1).alias("clnsig"),
            pl.col("pathogen_type").fill_null("-").alias("pathogen_type"),
        )
        # 注释tNGS/mNGS
        if not self.target_records.is_empty():
            for key in ["taxid", "species_taxid", "genus_taxid"]:
                records = records.join(self.target_records.rename({"taxid": key.replace("taxid", "target")}), how="left", left_on="taxid", right_on=key)
            records = records.with_columns(pl.col("taxid").fill_null(pl.col("species_target")).fill_null(pl.col("species_target")).fill_null(pl.lit("mNGS"))).drop(
                ["species_target", "genus_target"]
            )
        # 计算信号
        records = records.with_columns(
            pl.struct(["taxid", "rpm", "uniq_reads", "genus_ranking", "genus_rel_abund", "pathogen_type", "clnsig", "target"])
            .map_elements(lambda x: Pathogen.model_validate(x).calc_signal(self.sample_type).value, dtype=pl.String)
            .alias("signal")
        )
        return records


def set_options(sub_parser):
    parser = sub_parser.add_parser("tngs", help="Output tNGS result ")
    parser.add_argument("--pathogen_tsv", "-patho", type=Path, required=True, nargs="+", help="input, Pathogen TSV files")
    parser.add_argument("--reads_tsv", "-reads", type=Path, required=True, nargs="+", help="input, Reads TSV files")
    parser.add_argument("--resfinder_tsv", "-res", type=Path, required=True, nargs="+", help="input, ResFinder TSV file")
    parser.add_argument("--virulencefinder_tsv", "-vf", type=Path, required=True, nargs="+", help="input, VirulenceFinder TSV file")
    parser.add_argument("--target_tsv", "-tgt", type=Path, help="input, Target TSV file contains TaxID")
    parser.add_argument("--sample_type", "-type", choices=("呼吸道", "外周血", "脑脊液", "胸水", "骨髓", "腹水", "关节积液", "眼部"), help="input, Target TSV file contains TaxID")
    parser.add_argument("--ctrl_sample", "-ctrl", nargs="+", required=True, help="parameter, Control sample")
    parser.add_argument("--result_zip", "-out", type=Path, required=True, help="output, Result zip file")
    parser.set_defaults(func=lambda args: Result.model_validate(vars(args)).run())
